import sqlite3
import os
from pathlib import Path
import uuid
from datetime import datetime

# 获取项目根目录
ROOT_DIR = Path(__file__).parent.parent
DB_PATH = ROOT_DIR / "articles.db"

def init_database():
    """初始化数据库"""
    # 如果数据库文件已存在，先删除
    if os.path.exists(DB_PATH):
        os.remove(DB_PATH)
    
    # 创建数据库连接
    conn = sqlite3.connect(str(DB_PATH))
    cur = conn.cursor()
    
    # 创建用户表
    cur.execute("""
        CREATE TABLE IF NOT EXISTS users (
            id TEXT PRIMARY KEY,
            username TEXT NOT NULL UNIQUE,
            email TEXT NOT NULL UNIQUE,
            password_hash TEXT NOT NULL,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            last_login TIMESTAMP,
            status TEXT CHECK(status IN ('active', 'inactive', 'suspended')) DEFAULT 'active'
        )
    """)
    
    # 创建文章表
    cur.execute("""
        CREATE TABLE IF NOT EXISTS articles (
            id TEXT PRIMARY KEY,
            user_id TEXT NOT NULL,
            title TEXT NOT NULL,
            content TEXT NOT NULL,
            summary TEXT,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            status TEXT CHECK(status IN ('draft', 'published', 'archived')) DEFAULT 'draft',
            view_count INTEGER DEFAULT 0,
            FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE
        )
    """)
    
    # 创建文章参数表
    cur.execute("""
        CREATE TABLE IF NOT EXISTS article_params (
            id TEXT PRIMARY KEY,
            article_id TEXT NOT NULL UNIQUE,
            topic TEXT,
            style TEXT CHECK(style IN ('professional', 'casual', 'storytelling', 'technical', 'educational')) DEFAULT 'professional',
            length TEXT CHECK(length IN ('short', 'medium', 'long')) DEFAULT 'medium',
            audience TEXT CHECK(audience IN ('general', 'professional', 'young')) DEFAULT 'general',
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            FOREIGN KEY (article_id) REFERENCES articles(id) ON DELETE CASCADE
        )
    """)
    
    # 创建参考链接表
    cur.execute("""
        CREATE TABLE IF NOT EXISTS reference_urls (
            id TEXT PRIMARY KEY,
            article_id TEXT NOT NULL,
            url TEXT NOT NULL,
            title TEXT,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            FOREIGN KEY (article_id) REFERENCES articles(id) ON DELETE CASCADE
        )
    """)
    
    # 创建标签表
    cur.execute("""
        CREATE TABLE IF NOT EXISTS tags (
            id TEXT PRIMARY KEY,
            name TEXT NOT NULL UNIQUE,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    """)
    
    # 创建文章标签关联表
    cur.execute("""
        CREATE TABLE IF NOT EXISTS article_tags (
            article_id TEXT NOT NULL,
            tag_id TEXT NOT NULL,
            PRIMARY KEY (article_id, tag_id),
            FOREIGN KEY (article_id) REFERENCES articles(id) ON DELETE CASCADE,
            FOREIGN KEY (tag_id) REFERENCES tags(id) ON DELETE CASCADE
        )
    """)
    
    # 创建文章版本历史表
    cur.execute("""
        CREATE TABLE IF NOT EXISTS article_versions (
            id TEXT PRIMARY KEY,
            article_id TEXT NOT NULL,
            content TEXT NOT NULL,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            description TEXT,
            FOREIGN KEY (article_id) REFERENCES articles(id) ON DELETE CASCADE
        )
    """)
    
    # 创建AI模型表
    cur.execute('''
        CREATE TABLE IF NOT EXISTS ai_models (
            id TEXT PRIMARY KEY,
            name TEXT NOT NULL,
            provider TEXT NOT NULL,
            model_id TEXT NOT NULL,
            api_type TEXT NOT NULL,
            description TEXT,
            max_tokens INTEGER NOT NULL DEFAULT 4096,
            temperature_default REAL NOT NULL DEFAULT 0.7,
            api_key TEXT NOT NULL,
            api_endpoint TEXT,
            api_version TEXT,
            stream_enabled INTEGER NOT NULL DEFAULT 1,
            is_active INTEGER NOT NULL DEFAULT 1,
            created_at TIMESTAMP NOT NULL,
            updated_at TIMESTAMP NOT NULL
        )
    ''')
    
    # 创建提示词模板表
    cur.execute("""
        CREATE TABLE IF NOT EXISTS prompt_templates (
            id TEXT PRIMARY KEY,
            name TEXT NOT NULL,
            description TEXT,
            system_prompt TEXT,
            user_prompt TEXT NOT NULL,
            variables TEXT,  -- JSON 格式存储变量列表
            category TEXT,
            tags TEXT,  -- JSON 格式存储标签列表
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    """)
    
    # 创建索引
    cur.execute("CREATE INDEX IF NOT EXISTS idx_articles_user ON articles(user_id)")
    cur.execute("CREATE INDEX IF NOT EXISTS idx_articles_status ON articles(status)")
    cur.execute("CREATE INDEX IF NOT EXISTS idx_article_tags ON article_tags(tag_id)")
    cur.execute("CREATE INDEX IF NOT EXISTS idx_ref_urls_article ON reference_urls(article_id)")
    
    # 插入测试数据
    # 创建测试用户
    test_user_id = str(uuid.uuid4())
    cur.execute("""
        INSERT INTO users (id, username, email, password_hash)
        VALUES (?, ?, ?, ?)
    """, (test_user_id, 'test_user', 'test@example.com', 'hashed_password'))
    
    # 创建测试文章
    article_id = str(uuid.uuid4())
    cur.execute("""
        INSERT INTO articles (id, user_id, title, content, summary, status)
        VALUES (?, ?, ?, ?, ?, ?)
    """, (
        article_id,
        test_user_id,
        '测试文章',
        '这是一篇测试文章的内容...',
        '测试文章摘要',
        'draft'
    ))
    
    # 创建文章参数
    cur.execute("""
        INSERT INTO article_params (id, article_id, topic, style, length)
        VALUES (?, ?, ?, ?, ?)
    """, (str(uuid.uuid4()), article_id, '测试主题', 'professional', 'medium'))
    
    # 创建测试标签
    tag_id = str(uuid.uuid4())
    cur.execute("""
        INSERT INTO tags (id, name)
        VALUES (?, ?)
    """, (tag_id, '测试标签'))
    
    # 关联文章和标签
    cur.execute("""
        INSERT INTO article_tags (article_id, tag_id)
        VALUES (?, ?)
    """, (article_id, tag_id))
    
    # 提交更改
    conn.commit()
    conn.close()
    
    print(f"数据库初始化完成: {DB_PATH}")

if __name__ == "__main__":
    init_database() 